package com.csx.governance.traffic.adapter.inbound.webflux;

import com.csx.governance.traffic.core.context.TrafficContext;
import com.csx.governance.traffic.core.context.TrafficContextHolder;
import com.csx.governance.traffic.core.context.TrafficRequest;
import com.csx.governance.traffic.core.context.TrafficResponse;
import com.csx.governance.traffic.core.context.webflux.TrafficWebFluxRequest;
import com.csx.governance.traffic.core.context.webflux.TrafficWebFluxResponse;
import com.csx.governance.traffic.core.interceptor.TrafficHandlerInterceptorRunner;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;

import java.util.concurrent.atomic.AtomicReference;

/**
 * WebFlux 入栈，通过 {@link TrafficHandlerInterceptorRunner} 调用一系列 InboundHandlerInterceptor
 * @author cuisongxu
 * @date 2023/2/23 14:59
 */
public class TrafficSwimLaneWebFilter implements WebFilter {


    private final Logger logger = LoggerFactory.getLogger(TrafficSwimLaneWebFilter.class);

    private final TrafficHandlerInterceptorRunner handlerInterceptorRunner;

    public TrafficSwimLaneWebFilter(TrafficHandlerInterceptorRunner handlerInterceptorRunner) {
        this.handlerInterceptorRunner = handlerInterceptorRunner;
    }

    @Override
    public reactor.core.publisher.Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {

        TrafficWebFluxRequest trafficWebFluxRequest = new TrafficWebFluxRequest(exchange);
        TrafficWebFluxResponse trafficWebFluxResponse = new TrafficWebFluxResponse(exchange);
        AtomicReference<Throwable> ex = new AtomicReference<>();
        // 将 请求对象 和 响应对象 保存在上下文中
        buildContext(trafficWebFluxRequest, trafficWebFluxResponse);
        try {
            handlerInterceptorRunner.preHandle(trafficWebFluxRequest, trafficWebFluxResponse);
        }catch (Exception e) {
            handlerInterceptorRunner.afterCompletion(trafficWebFluxRequest, trafficWebFluxResponse, e);
            return Mono.error(e);
        }
        return chain.filter(exchange)
                .doOnSuccess(result -> handlerInterceptorRunner.postHandle(trafficWebFluxRequest, trafficWebFluxResponse))
                .doOnError(e -> {
                    ex.set(e);
                    logger.error("TrafficSwimLaneWebFilter doFilter error because: {}", e.getMessage());
                })
                .doFinally(signalType -> handlerInterceptorRunner.afterCompletion(trafficWebFluxRequest, trafficWebFluxResponse, ex.get()));
    }

    private void buildContext(TrafficRequest request, TrafficResponse response) {
        TrafficContext trafficContext = new TrafficContext(request, response);
        TrafficContextHolder.setContext(trafficContext);
    }
}
